# -*- coding: utf-8 -*-

"""
    2025-01-30 by Jie Zheng
    from: 2017-02-03, 2016M nd others, Jie Zheng
    Some wcs routines from idl
    NOTE: NOT ALL copy, simplified for our purpose only

    alpha,delta <-crval-> phi,theta <--> xsi,eta <-CD-> u,v <-(crpix)-> x,y
    x,y - crpix = u,v                     | u,v + crpix = x,y
    u,v x CD = xsi,eta                    | xsi,eta x RCD = u,v
    xsi,eta convert = phi,theta           | phi,theta convert = xsi,eta
    phi,theta rotate(crval) = alpha,delta | alpha,delta rotate(crval) = phi,theta
"""


import numpy as np


def cd_inv(cd: np.ndarray):
    """
    Inverse CD matrix
    :param cd: CD matrix
    :return: inversed CD
    """
    v = cd[0, 0] * cd[1, 1] - cd[0, 1] * cd[1, 0]
    i = np.array([[cd[1, 1], -cd[0, 1]], [-cd[1, 0], cd[0, 0]]]) / v
    return i


def xe2pt(xsi:np.ndarray, eta:np.ndarray):
    """
    Convert (xsi,eta) coordinates to spherical coordinates (phi,theta)
    center at north pole
    :param xsi: xsi
    :param eta: eta
    :return: theta,phi
    """
    r = np.sqrt(xsi * xsi + eta * eta)
    theta = np.rad2deg(np.arctan2(1, np.deg2rad(r)))
    phi = np.rad2deg(np.arctan2(xsi, -eta))
    return phi, theta


def pt2xe(phi:np.ndarray, theta:np.ndarray):
    """
    Convert spherical coordinates (phi,theta) to xsi,eta system
    :param phi: phi
    :param theta: theta
    :return: xsi,eta
    """
    xsi = np.empty_like(phi, dtype=float) + np.nan
    eta = np.empty_like(phi, dtype=float) + np.nan

    g = np.where(theta > 0)[0]
    if len(g) > 0:
        r_theta = 1.0 / np.deg2rad(np.tan(np.deg2rad(theta[g])))
        xsi[g] = r_theta * np.sin(np.deg2rad(phi[g]))
        eta[g] = - r_theta * np.cos(np.deg2rad(phi[g]))
    return xsi, eta


def xy2uv(x:np.ndarray, y:np.ndarray, crpix:np.ndarray|list|tuple):
    """
    x,y to u,v
    """
    return x - crpix[0], y - crpix[1]


def uv2xy(u:np.ndarray, v:np.ndarray, crpix:np.ndarray|list|tuple):
    """
    u,v to x,y
    """
    return u + crpix[0], v + crpix[1]


def pt2ad(phi:np.ndarray, theta:np.ndarray, crval:np.ndarray|list|tuple):
    """
    phi/theta to alpha,delta
    """
    alpha, delta = _wcsrotate_(crval, phi, theta, "S")
    return alpha, delta


def ad2pt(alpha:np.ndarray, delta:np.ndarray, crval:np.ndarray|list|tuple):
    """
    alpha, delta to phi, theta
    """
    phi, theta = _wcsrotate_(crval, alpha, delta, "N")
    return phi, theta


def uv2xe(u:np.ndarray, v:np.ndarray, cd:np.ndarray):
    """
    u,v to xsi,eta
    """
    xsi = cd[0,0] * u + cd[0,1] * v
    eta = cd[1,0] * u + cd[1,1] * v
    return xsi, eta


def xe2uv(xsi:np.ndarray, eta:np.ndarray, cd:np.ndarray):
    """
    xsi,eta to u,v
    """
    # CD -> CD^-1
    rcd = cd_inv(cd)
    u = xsi * rcd[0, 0] + eta * rcd[0, 1]
    v = xsi * rcd[1, 0] + eta * rcd[1, 1]
    return u, v


def xe2ad(xsi:np.ndarray, eta:np.ndarray, crval:np.ndarray|list|tuple):
    """
    xsi,eta to alpha,delta, througn phi,theta
    """
    phi, theta = xe2pt(xsi, eta)
    alpha, delta = pt2ad(phi, theta, crval)
    return alpha, delta


def ad2xe(alpha:np.ndarray, delta:np.ndarray, crval:np.ndarray|list|tuple):
    """
    alpha,delta to xsi,eta, througn phi,theta
    """
    phi, theta = ad2pt(alpha, delta, crval)
    xsi, eta = pt2xe(phi, theta)
    return xsi, eta


def xy2ad(x:np.ndarray, y:np.ndarray, crpix:np.ndarray|list|tuple, cd:np.ndarray, crval:np.ndarray|list|tuple):
    """
    x,y to alpha, delta, through u,v, xsi,eta, phi,theta
    """
    u, v = xy2uv(x, y, crpix, cd)
    xsi, eta = uv2xe(u, v, cd)
    phi, theta = xe2pt(xsi, eta)
    alpha, delta = pt2ad(phi, theta, crval)
    return alpha, delta


def ad2xy(alpha:np.ndarray, delta:np.ndarray, crpix:np.ndarray|list|tuple, cd:np.ndarray, crval:np.ndarray|list|tuple):
    """
    alpha, delta to x,y, through phi,theta, xsi,eta, u,v
    """
    phi, theta = ad2pt(alpha, delta, crval)
    xsi, eta = pt2xe(phi, theta)
    u, v = xe2uv(xsi, eta, cd)
    x, y = uv2xy(u, v, crpix)
    return x, y


def _wcsrotate_(crval, c01, c02, dir):
    """
    Rotate between standard (e.g. celestial) and native coordinates
    Simplified for our purpose
    Reference from wcs_rotate of astron of IDL
    :param crval: crval of wcs info, 2-tuple/list/array
    :param c01: first coord, c11 if N, phi if S
    :param c02: second coord, dec if N, theta if S
    :param dir: direction,
                N, to native, ra,dec -> phi,theta, (IDL positive dir)
                S, to standard, phi,theta -> ra,dec (IDL reverse dir)
    :return: c11, c12
    """
    isarray = isinstance(c01, np.ndarray)
    # original from PV1, but in this case, they are constants
    phi0 = 0.0
    theta0 = 90.0
    longpole = 180.0
    latpole = 90.0
    phi_p = np.deg2rad(longpole)
    sp, cp = 0.0, -1.0
    alpha_p = np.deg2rad(crval[0])
    delta_p = np.deg2rad(crval[1])
    at_pole = crval[1] >= 90 - 1e-12  # pole and a minor error
    if at_pole and (alpha_p - phi_p) % (2 * np.pi) == np.pi:
        # Native and celestial frames coincide, No rotation needed
        c11, c12 = c01.copy(), c02.copy()
    else:
        sa = np.sin(alpha_p)
        ca = np.cos(alpha_p)
        sd = np.sin(delta_p)
        cd = np.cos(delta_p)
        if at_pole : cd = 0.0  # suppress rounding errors
        # calculate rotation matrix
        r = np.array([[-sa * sp - ca * cp * sd,  ca * sp - sa * cp * sd, cp * cd],
                      [ sa * cp - ca * sp * sd, -ca * cp - sa * sp * sd, sp * cd],
                      [      ca * cd,                 sa * cd,              sd  ]])
        if dir.upper() == "S":
            c11, c12 = c02.copy(), c01.copy()  # NOTE: this is not same as upper
            g = np.where(np.isfinite(c01) & np.isfinite(c02))[0]
            if len(g) > 0:
                if isarray:
                    phi1 = np.deg2rad(c01[g])
                    theta1 = np.deg2rad(c02[g])
                else:
                    phi1 = np.deg2rad(c01)
                    theta1 = np.deg2rad(c02)
            else:
                theta1, phi1 = 0.0, 0.0   # suppress warning
        else:
            if isarray:
                c11 = c01.copy()
                c12 = c02.copy()  # add this for code checking
                phi1 = c11.copy()
                g = np.where(np.abs(c01) != 180.0)  # Avoid roundoff error
                if len(g) > 0:
                    phi1[g] = np.deg2rad(c01[g])
            else:
                c11 = c01
                c12 = c02  # add this for code checking
                phi1 = c11
                g = np.where(np.abs(c01) != 180.0)  # Avoid roundoff error
                if len(g) > 0:
                    phi1 = np.deg2rad(c01[g])
            r = r.T  # since array order different between IDL and python, r.T move here
            theta1 = np.deg2rad(c02)
        l = np.cos(theta1) * np.cos(phi1)
        m = np.cos(theta1) * np.sin(phi1)
        n = np.sin(theta1)
        b0 = r[0, 0] * l + r[1, 0] * m + r[2, 0] * n
        b1 = r[0, 1] * l + r[1, 1] * m + r[2, 1] * n
        b2 = r[0, 2] * l + r[1, 2] * m + r[2, 2] * n
        # Account for possible roundoff
        if isarray:
            b2[np.where(b2 < -1)[0]] = -1.0
            b2[np.where(b2 > 1)[0]] = 1.0
        else:
            if b2 < -1: b2 = -1.0
            if b2 > 1: b2 = 1.0
        if isarray and dir.upper() == "S":
            c11[g] = np.rad2deg(np.arctan2(b1, b0))
            c12[g] = np.rad2deg(np.arcsin(b2))
        else:
            c11 = np.rad2deg(np.arctan2(b1, b0))
            c12 = np.rad2deg(np.arcsin(b2))
    # if dir.upper() == "N":
    ix = np.where(c11 < 0)[0]
    if len(ix) > 0:
        if isarray:
            c11[ix] += 360.0
        else:
            c11 += 360.0
    return c11, c12
